import os
import time
import sqlite3
import datetime
import baostock as bs
import pandas as pd
import numpy as np
import config
from tqdm import tqdm
from utils import download_db_from_cos, upload_db_to_cos

# 常量配置
HISTORY_DAYS = 400          # 需要获取的历史自然日天数
SLEEP_TIME = 0.1            # 请求间隔(秒)
TARGET_TRADING_DAYS = 240   # 目标交易日数

def get_trade_dates_cache(start_date, end_date):
    """获取交易日历缓存"""
    rs = bs.query_trade_dates(start_date, end_date)
    trade_dates = []
    if rs.error_code == '0':
        while rs.next():
            trading_date, is_trading_day = rs.get_row_data()
            if is_trading_day == '1':
                trade_dates.append(trading_date)
    return set(trade_dates) if trade_dates else set()

def next_trading_day(last_date, trade_dates_set):
    """获取下一个交易日"""
    if not last_date or not trade_dates_set:
        return None
    
    # 转换为datetime格式
    last_dt = datetime.datetime.strptime(last_date, "%Y-%m-%d")
    
    # 检查后续1-5天
    for delta in range(1, 6):
        check_date = (last_dt + datetime.timedelta(days=delta)).strftime("%Y-%m-%d")
        if check_date in trade_dates_set:
            return check_date
    return None

def get_valid_stocks():
    """获取符合条件的股票列表"""
    # 只保留主板和创业板
    valid_prefixes = [
        'sh.600', 'sh.601', 'sh.603', 'sh.605', 
        'sz.000', 'sz.001', 'sz.002', 'sz.003', 
        'sz.300'
    ]
    
    stock_list = []
    today_str = datetime.datetime.now().strftime("%Y-%m-%d")
    min_ipo_date = (datetime.datetime.now() - datetime.timedelta(days=180)).strftime("%Y-%m-%d")
    
    try:
        rs = bs.query_stock_basic()
        if rs.error_code != '0':
            return []
            
        while rs.next():
            code, name, ipo_date, out_date, status = rs.get_row_data()[0:5]
            
            # 有效性检查
            if not any(code.startswith(p) for p in valid_prefixes):
                continue
                
            if status != '1':  # 非上市状态
                continue
                
            # 排除ST股票
            if 'ST' in name or '*ST' in name:
                continue
                
            # 排除科创板
            if code.startswith('sh.688'):
                continue
                
            if out_date and out_date < today_str:  # 已退市
                continue
                
            # 排除上市不足半年的新股
            if ipo_date and ipo_date > min_ipo_date:
                continue
                
            stock_list.append(code)
                
    except Exception as e:
        print(f"获取股票列表异常: {e}")
        
    return stock_list

def get_stock_data(stock_code, start_date, end_date):
    """获取单只股票数据"""
    if start_date > end_date:
        return None, "开始日期大于结束日期"
    
    max_retries = 2
    last_error = ""
    
    for attempt in range(max_retries):
        try:
            # 创业坂使用后复权，其他使用前复权
            adjust = "2" if stock_code.startswith('sz.300') else "3"
            
            rs = bs.query_history_k_data_plus(
                stock_code,
                "date,code,open,high,low,close,volume,amount,adjustflag",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag=adjust
            )
            
            if rs.error_code != '0':
                last_error = f"[{rs.error_code}] {rs.error_msg}"
                time.sleep(1)
                continue
                
            # 处理数据
            data_list = []
            while rs.next():
                data_list.append(rs.get_row_data())
                
            if not data_list:
                return None, "无有效数据"
            
            # 转换为DataFrame
            df = pd.DataFrame(data_list, columns=rs.fields)
            
            # 过滤有效日期格式
            df = df[df['date'].str.len() == 10]
            return df, None
            
        except Exception as e:
            last_error = f"内部错误: {e}"
            time.sleep(1)
            
    return None, last_error or "获取失败"

def save_to_database(conn, df):
    """保存数据到数据库"""
    if df is None or df.empty:
        return False
    
    try:
        cursor = conn.cursor()
        insert_data = []
        for _, row in df.iterrows():
            insert_data.append((
                row['code'], row['date'], 
                row.get('open', 0), row.get('high', 0), 
                row.get('low', 0), row.get('close', 0),
                row.get('volume', 0), row.get('amount', 0), 
                row.get('adjustflag', 0)
            ))
        
        # 批量插入（忽略重复）
        cursor.executemany('''
            INSERT OR IGNORE INTO stock_data 
                (code, date, open, high, low, close, volume, amount, adjustflag)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', insert_data)
        
        conn.commit()
        return True
    
    except Exception as e:
        print(f"保存错误: {str(e)}")
        return False

def get_stock_status(conn, stock_codes, start_date, end_date):
    """获取股票状态信息"""
    if not stock_codes:
        return {}
    
    stock_info = {}
    
    try:
        # 查询现有股票数据状态
        placeholders = ','.join(['?'] * len(stock_codes))
        query = f"""
            SELECT code, MAX(date) as max_date, COUNT(*) as count 
            FROM stock_data 
            WHERE code IN ({placeholders}) 
              AND date BETWEEN ? AND ?
            GROUP BY code
        """
        params = tuple(stock_codes) + (start_date, end_date)
        
        cursor = conn.cursor()
        cursor.execute(query, params)
        db_results = {}
        for row in cursor.fetchall():
            code, max_date, count = row
            db_results[code] = {'max_date': max_date, 'count': count}
        
        # 获取IPO信息
        rs = bs.query_stock_basic()
        stock_basics = {}
        if rs.error_code == '0':
            while rs.next():
                row_data = rs.get_row_data()
                if len(row_data) >= 3:
                    stock_basics[row_data[0]] = row_data[2]
        
        # 合并结果
        for code in stock_codes:
            db_data = db_results.get(code, {'max_date': None, 'count': 0})
            stock_info[code] = {
                'max_date': db_data['max_date'],
                'count': db_data['count'],
                'ipo_date': stock_basics.get(code, None)
            }
    
    except Exception as e:
        print(f"获取股票状态时出错: {str(e)}")
        for code in stock_codes:
            stock_info[code] = {
                'max_date': None,
                'count': 0,
                'ipo_date': None
            }
    
    return stock_info

def main_download():
    print("===== 股票数据下载任务开始 =====")
    
    # 1. 下载数据库
    print("从COS下载数据库...")
    download_db_from_cos()
    
    # 获取数据库路径
    cfg = config.COS_CONFIG
    db_config = config.DB_FILES
    stock_data_db = os.path.join(cfg["DB_DIR"], db_config["STOCK_DATA"])
    
    # 2. 登录baostock
    print("登录baostock系统...")
    lg = bs.login()
    if lg.error_code != '0':
        error_msg = f"登录失败: {lg.error_msg}"
        print(error_msg)
        send_wechat_notification("数据下载失败", error_msg)
        return
    
    # 3. 设置日期范围
    today = datetime.datetime.now()
    end_date = today.strftime("%Y-%m-%d")
    start_date = (today - datetime.timedelta(days=HISTORY_DAYS)).strftime("%Y-%m-%d")
    
    # 4. 获取交易日历
    print("获取交易日历...")
    trade_dates_cache = get_trade_dates_cache(start_date, end_date)
    
    # 5. 连接数据库
    conn = sqlite3.connect(stock_data_db)
    cursor = conn.cursor()
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS stock_data (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        code TEXT NOT NULL,
        date TEXT NOT NULL,
        open REAL,
        high REAL,
        low REAL,
        close REAL,
        volume INTEGER,
        amount REAL,
        adjustflag INTEGER,
        UNIQUE(code, date)
    )
    ''')
    conn.commit()
    
    # 6. 获取有效股票列表
    print("获取有效股票列表...")
    stock_codes = get_valid_stocks()
    if not stock_codes:
        print("未找到有效股票")
        return
    
    status = get_stock_status(conn, stock_codes, start_date, end_date)
    
    # 7. 数据下载
    total_count = len(status)
    success_count = 0
    skip_count = 0
    fail_count = 0
    download_days = 0
    
    print(f"开始下载 {total_count} 只股票数据...")
    for code, info in tqdm(status.items(), total=total_count):
        need_update = False
        update_reason = "无需更新"
        actual_start_date = start_date
        
        # 检查是否需要更新
        if info['count'] == 0:  # 首次下载
            need_update = True
            update_reason = "首次下载"
            if info['ipo_date'] and info['ipo_date'] > start_date:
                actual_start_date = info['ipo_date']
                
        elif info['max_date']:  # 已有数据
            next_date = next_trading_day(info['max_date'], trade_dates_cache)
            if next_date and next_date <= end_date:
                need_update = True
                update_reason = "有新交易日"
                actual_start_date = next_date
        
        # 执行下载
        if need_update:
            df, error = get_stock_data(code, actual_start_date, end_date)
            if df is None or df.empty:
                fail_count += 1
                print(f"  {code} 下载失败: {error}")
            else:
                if save_to_database(conn, df):
                    success_count += 1
                    download_days += len(df)
                    print(f"  {code} 新增 {len(df)} 条数据")
                else:
                    fail_count += 1
                    print(f"  {code} 数据保存失败")
        else:
            skip_count += 1
    
    # 8. 清理
    conn.close()
    bs.logout()
    
    # 9. 上传更新后的数据库
    print("上传更新后的数据库到COS...")
    upload_db_to_cos()
    
    # 10. 发送通知
    summary = (
        f"股票数据下载完成！\n"
        f"股票总数: {total_count}\n"
        f"成功下载: {success_count}\n"
        f"跳过: {skip_count}\n"
        f"失败: {fail_count}\n"
        f"新增数据: {download_days} 条"
    )
    
    print(summary)
    send_wechat_notification("股票数据下载报告", summary)
    
    print("===== 数据下载任务完成 =====")

if __name__ == '__main__':
    main_download()
